import tlsh


class MyTlsh:

    def modify_res(self, res):
        """
        将tlsh的score 转换成0到1之间
        根据推荐的指标100为临界点--》0.5
        100的正好为0.5 ，最大值=1，最小值等于0（样本量较少的话，可能会有bug）
        """
        a = max(res)
        b = min(res)
        temp = [0] * len(res)
        for i in range(len(res)):
            b_min = max(b, 100)
            a_max = min(a, 100)
            if res[i] > 100:
                tar = (1 - (res[i] - b_min) / (a - b_min)) / 2
            else:
                tar = (1 - (res[i] - b) / (a_max - b)) / 2 + 0.5
            temp[i] = tar
        return temp

    def score(self, a, b):
        """
        输入:对比的两个样本tlsh
        输出:这两个样本的tlsh得分
        这里设置的推荐阈值是100
        score越低--->相似度越高
        """
        res = []
        for i in range(len(a)):
            score = tlsh.diff(a[i], b[i])
            res.append(score)
        return self.modify_res(res)

    def calculateTlsh(self, samplePath):
        """
        input: sample path
        output: tlse value
        """
        return tlsh.hash(open(samplePath, "rb").read())

    def getTopkscore(self, tartlsh, allTlsh, k):
        """
        tarTlsh: the tlsh value of target sample (id, tlsh)
        allTlsh: the collections of all tlsh value save in database [(id1,tlsh),(id2,tlsh)...]
        output: top k tlsh values of similarity score
        tip: id is md5 value
        """
        tarId = tartlsh[0]
        tarTlsh = tartlsh[1]
        res = []
        for dataId, dataTlsh in allTlsh:
            if dataId == tarId:
                continue
            score = tlsh.diff(tarTlsh, dataTlsh)
            res.append((dataId, score))
        top_k = sorted(res, key=lambda x: x[1])[:k]
        print("粗筛结果为:" + str(top_k))
        return [item[0] for item in top_k]
